//
// Created by EmFi on 2023/5/15.
//

#ifndef DBMS_BPLUSTREE_H
#define DBMS_BPLUSTREE_H

#include <queue>
#include <list>
#include "InternalNode.h"
#include "LeafNode.h"

template <class K, class V>
class BPlusTree {
protected:
    // 树的阶
    const int order;
    // 根节点
    Node<K> *root;
private:
    // 查找 key 对应的叶子结点，如果没找到返回 NULL
    LeafNode<K, V> *findLeafNode(K key) {
        Node<K> *node = root;
        while (node != nullptr) {
            // 如果当前结点为叶子结点，跳出循环
            if (node->getType() == NodeType::LEAF) {
                break;
            }

            // 找到第一个键值大于等于 key 的位置
            int index = 0;
            while (key >= node->getKey(index) && index < node->getLength()) {
                index++;
            }

            node = ((InternalNode<K>*) node)->getPointer(index);
        }

        return (LeafNode<K, V>*) node;
    }

    // 在内部节点 node 插入 key，并设置 key 的右指针为 rightSon
    // rightSon 既可以是内部节点也可以是叶子结点
    bool insertInternalNode(InternalNode<K> *node, K key, Node<K> *rightSon) {
        if (node == nullptr || node->getType() == NodeType::LEAF) {
            return false;
        }

        // 判断当前内部节点状态
        if (node->getLength() < node->getCapacity()) {
            // 如果未满，直接插入
            return node->insert(key, rightSon);
        }

        // 如果当前内部节点已满，分裂内部节点
        auto *newInternal = new InternalNode<K>(this->order);
        K newKey = node->split(newInternal, key);

        // 根据分裂后两个内部节点的长度关系判断 key 插入的位置
        if (node->getLength() < newInternal->getLength()) {
            node->insert(key, rightSon);
        } else if(node->getLength() > newInternal->getLength()) {
            newInternal->insert(key, rightSon);
        } else {
            // 两者相等，即 key 在第 V 和 V+1 个键中间，把 rightSon 挂到新结点的第 0 个指针上
            newInternal->setPointer(0, rightSon);
            rightSon->setFather(newInternal);
        }

        // 检查父节点
        auto *father = (InternalNode<K> *) node->getFather();
        if (father == nullptr) {
            // 如果当前内部节点无父节点，说明该内部节点原来是根节点
            // 新建根节点，并将 node 作为根节点的第 0 个指针，newInternal 作为第 1 个指针
            auto *newRoot = new InternalNode<K>(this->order);
            newRoot->setKey(0, newKey);
            newRoot->setPointer(0, node);
            newRoot->setPointer(1, newInternal);
            newRoot->setLength(1);

            // 处理父节点关系
            node->setFather(newRoot);
            newInternal->setFather(newRoot);

            this->root = newRoot;
            return true;
        }

        // 如果当前内部节点有父节点，继续在父节点插入 newKey，并设置 newKey 的右指针为 newInternal
        return this->insertInternalNode(father, newKey, newInternal);
    }

    bool removeInternalNode(InternalNode<K> *node, K key) {
        // 删除键，如果失败一定是没有找到，直接返回失败
        if (!node->remove(key)) {
            return false;
        }

        // 处理父结点
        auto *father = (InternalNode<K> *) (node->getFather());
        if (father == nullptr) {
            // 如果一个数据都没有了，把根结点的第 0 个指针作为新的根结点
            if (node->getLength() == 0) {
                this->root = node->getPointer(0);

                // 重置新根节点的父节点为空
                this->root->setFather(nullptr);

                delete node;
            }
            return true;
        }

        // 删除后结点填充度仍 >=50%
        if (node->getLength() >= this->order) {
            // 重置父节点中该位置的键为 node 的第 0 个键
            for (int i = 0; i < father->getLength(); i++) {
                if (father->getKey(i) == key) {
                    father->setKey(i, node->getKey(0));
                    break;
                }
            }
            return true;
        }

        // 删除后结点填充度 <50%

        // 找到最近的一个兄弟结点，优先找左兄弟
        auto *brother = (InternalNode<K> *) node->getPrevBrother();
        bool isNextBro = false;
        if (brother == nullptr) {
            brother = (InternalNode<K> *) node->getNextBrother();
            isNextBro = true;
        }

        // 如果兄弟结点填充度 >50%，则向兄弟节点借一个节点（仅指针，键应设置为借的指针上节点的第 0 个键）插入到本节点
        if (brother->getLength() > this->order) {
            // 将兄弟结点的某一个指针移动到本结点，用借完之后较大节点的 key 替换父结点中 node 和 brother 中间的 key
            if (!isNextBro) {
                // 如果是左兄弟

                // 右旋
                for (int i = 1; i <= father->getLength(); i++) {
                    if (father->getPointer(i) == node) {
                        // 将根节点上 brother 左边的键以及兄弟的最后一个节点的右指针插入到本节点的第 0 个位置，且指针的位置在插入位置左侧
                        node->insert(father->getKey(i - 1), brother->getPointer(node->getLength()), false);
                        // 将父节点上 node 左边的键设置成 brother 的最后一个键
                        father->setKey(i - 1, brother->getKey(node->getLength() - 1));
                    }
                }
                // 删除 brother 的最后一个键以及其右指针
                brother->remove(brother->getKey(node->getLength() - 1));
            } else {
                // 如果是右兄弟

                // 左旋
                for (int i = 1; i <= father->getLength(); i++) {
                    if (father->getPointer(i) == brother) {
                        // 将根节点上 brother 左边的键以及兄弟的第 0 个节点的左指针插入到本节点的最后一个位置，且指针的位置在插入位置右侧
                        node->insert(father->getKey(i - 1), brother->getPointer(0));
                        // 将父节点上 brother 左边的键设置成 brother 的第 0 个键
                        father->setKey(i - 1, brother->getKey(0));
                    }
                }
                // 删除 brother 的第 0 个键及其左指针
                brother->remove(brother->getKey(0), false);
            }

            return true;
        }

        // 如果兄弟结点填充度 =50%，则需要把本结点与兄弟结点进行合并（合并到较小的节点上，这样父节点中要删除的就一定是键的右指针）
        K delKey;
        if (!isNextBro) {
            brother->merge(node);

            // 要删除的 key 为父节点中 node 和 brother 中间的键
            int nodeIndex, brotherIndex;
            for (int i = 0; i < father->getPointerCapacity(); i++) {
                if (father->getPointer(i) == node) {
                    nodeIndex = i;
                }
                if (father->getPointer(i) == brother) {
                    brotherIndex = i;
                }
            }
            delKey = father->getKey(min(nodeIndex, brotherIndex));

            delete node;
        } else {
            node->merge(brother);

            // 要删除的 key 为父节点中 node 和 brother 中间的键
            int nodeIndex, brotherIndex;
            for (int i = 0; i < father->getPointerCapacity(); i++) {
                if (father->getPointer(i) == node) {
                    nodeIndex = i;
                }
                if (father->getPointer(i) == brother) {
                    brotherIndex = i;
                }
            }
            delKey = father->getKey(min(nodeIndex, brotherIndex));

            delete brother;
        }

        return this->removeInternalNode(father, delKey);
    }
public:
    explicit BPlusTree(int order): order(order), root(nullptr) {}

    virtual ~BPlusTree() {
        delete this->root;
    }

    virtual void clear() {
        delete this->root;
        this->root = nullptr;
    }

    // 查找 key 对应的 value，如果没有找到返回 NULL
    V find(K key) {
        if (key == NULL) {
            return NULL;
        }

        LeafNode<K, V> *node = this->findLeafNode(key);

        if (node == nullptr) {
            return NULL;
        }

        // 在叶子结点中继续寻找
        return node->find(key);
    }

    // 查找 key 在 [start, end] 范围内的 value 列表
    list<V> find(K start, K end) {
        list<V> list;

        if (this->root == nullptr || start == NULL || end == NULL) {
            return list;
        }

        // 查找 start 对应的叶子结点
        LeafNode<K, V> *node = this->findLeafNode(start);

        // 第一个节点标识
        // node 中不一定包含 start，在这种情况下，需要的键在下一个兄弟中
        bool flag = true;
        // 从 start 对应的叶子节点开始，往后面的叶子节点查找 end
        while (node != nullptr || flag) {
            // node 中是否有需要的键
            bool found = false;

            // 在 node 中查找满足要求的键
            for (int i = 0; i < node->getLength(); i++) {
                K k = node->getKey(i);
                if (k >= start && k <= end) {
                    found = true;
                    list.push_back(node->getValue(i));
                }
            }

            // 如果不是第一个节点，本节点中没有找到需要的键，则退出循环
            if (!flag && !found) {
                break;
            }

            flag = false;
            node = node->getNextNode();
        }

        return list;
    }

    // 判断树中是否存在某 key
    bool contains(K key) {
        if (key == NULL) {
            return false;
        }
        return this->find(key) != NULL;
    }

    bool insert(K key, V val) {
        if (key == NULL) {
            return false;
        }

        // 查找欲插入的叶子结点
        LeafNode<K, V> *node = this->findLeafNode(key);

        if (node != nullptr) {
            // 如果存在该键，则不进行插入
            if (node->find(key) != NULL) {
                return false;
            }
        } else {
            // 如果没有找到，说明整个树是空的
            node = new LeafNode<K, V>(this->order);
            this->root = node;
        }

        // 判断当前叶子节点状态
        if (node->getLength() < node->getCapacity()) {
            // Case 1. 当前叶子结点未满
            // 直接在该叶子插入
            return node->insert(key, val);
        }

        // 如果当前叶子结点已满，将原叶子结点后一半数据剪切到新结点
        auto *newLeaf = new LeafNode<K, V>(this->order);
        K newKey = node->split(newLeaf);

        // 根据 key 判断是插入到原结点还是新结点中
        if (key < newKey) {
            node->insert(key, val);
        } else {
            newLeaf->insert(key, val);
        }

        // 检查父节点
        auto *father = (InternalNode<K> *) node->getFather();
        if (father == nullptr) {
            // 如果当前叶子节点无父节点，说明该叶子节点原来是根节点

            // Case 2. 当前叶子结点已满，且当前叶子节点为根节点
            // 首先把叶子节点分裂（已经完成，新节点为 newLeaf），然后新建根节点，并将 node 作为根节点的第 0 个指针，newLeaf 作为第 1 个指针

            auto *newRoot = new InternalNode<K>(this->order);
            newRoot->setKey(0, newKey);
            newRoot->setPointer(0, node);
            newRoot->setPointer(1, newLeaf);
            newRoot->setLength(1);

            // 处理父节点关系
            node->setFather(newRoot);
            newLeaf->setFather(newRoot);

            this->root = newRoot;
            return true;
        }

        // Case 3. 当前叶子结点已满，但其父结点未满
        // 首先把叶子结点分裂（已经完成，新节点为 newLeaf），然后选择插入原结点或新结点，再修改父结点的指针

        // Case 4. 当前叶子结点已满，且其父结点已满
        // 首先把叶子结点分裂（已经完成，新节点为 newLeaf），然后选择插入原结点或新结点；接着把父结点分裂，再修改祖父结点的指针（递归）

        return this->insertInternalNode(father, newKey, newLeaf);
    }

    bool update(K key, V val) {
        if (key == NULL) {
            return false;
        }

        // 查找欲更新的叶子结点
        LeafNode<K, V> *node = this->findLeafNode(key);

        if (node == nullptr) {
            return false;
        }

        for (int i = 0; i < node->getLength(); i++) {
            if (node->getKey(i) == key) {
                node->setValue(i, val);
                return true;
            }
        }

        return false;
    }

    bool remove(K key) {
        if (key == NULL) {
            return false;
        }

        // 查找欲删除的叶子结点
        LeafNode<K, V> *node = this->findLeafNode(key);

        if (node == nullptr) {
            return false;
        }

        // 在叶子节点中删除数据，如果失败一定是叶子节点中没有该键
        if (!node->remove(key)) {
            return false;
        }

        // 调整其他节点关系

        // 处理父结点
        auto *father = (InternalNode<K> *) (node->getFather());
        if (father == nullptr) {
            // 如果当前叶子结点无父节点，说明该叶子节点原来是根节点

            // Case 1. 如果当前删除的是根（叶子)结点中的键，直接删除（已经完成）

            // 如果根结点为空，删除根结点
            if (node->getLength() == 0) {
                delete node;

                this->root = nullptr;
            }
            return true;
        }

        // 如果当前叶子节点有父节点

        // 删除后叶子结点填充度仍 >=50%
        if (node->getLength() >= this->order) {
            // Case 2. 如果删除键的是父节点中的键，删除后该叶子结点填充度仍 >=50%
            // 重置父节点中该位置的键为叶子结点的第 0 个键
            for (int i = 0; i < father->getLength(); i++) {
                if (father->getKey(i) == key) {
                    father->setKey(i, node->getKey(0));
                    break;
                }
            }
            return true;
        }

        // 删除后叶子结点填充度 <50%

        // 找到最近的一个兄弟结点（必须调用 getBrother，不能调用叶子节点的 getNode，避免兄弟节点是其他树上的节点），优先找左兄弟
        auto *brother = (LeafNode<K, V> *) node->getPrevBrother();
        bool isNextBro = false;
        if (brother == nullptr) {
            brother = (LeafNode<K, V> *) node->getNextBrother();
            isNextBro = true;
        }

        // 如果兄弟结点填充度 >50%，则向兄弟节点借一个节点插入到本节点
        if (brother->getLength() > this->order) {
            // 将兄弟结点的某一个指针移动到本结点，用借完之后较大节点的 key 替换父结点中 node 和 brother 中间的 key
            if (!isNextBro) {
                // 如果是左兄弟

                // 移动左兄弟的最后一个节点到本结点
                node->insert(brother->getKey(brother->getLength() - 1), brother->getValue(brother->getLength() - 1));
                brother->remove(brother->getKey(brother->getLength() - 1));

                // 将父节点上 node 左侧的 key 设置成 node 的第 0 个键
                for (int i = 1; i <= father->getLength(); i++) {
                    if (father->getPointer(i) == node) {
                        father->setKey(i - 1, node->getKey(0));
                    }
                }
            } else {
                // 如果是右兄弟

                // 移动右兄弟的第 0 个节点到本结点
                node->insert(brother->getKey(0), brother->getValue(0));
                brother->remove(brother->getKey(0));

                // 将父节点上 brother 左侧的 key 设置成 brother 的第 0 个键
                for (int i = 1; i <= father->getLength(); i++) {
                    if (father->getPointer(i) == brother) {
                        father->setKey(i - 1, brother->getKey(0));
                    }
                }
            }

            return true;
        }

        // 如果兄弟结点填充度 =50%，则需要把本结点与兄弟结点进行合并（合并到较小的节点上，这样父节点中要删除的就一定是键的右指针）
        K delKey;
        if (!isNextBro) {
            brother->merge(node);

            // 处理叶子节点关系，此时需要删除的是 brother 的右兄弟
            brother->setNextNode(node->getNextNode());
            if (node->getNextNode()) {
                node->getNextNode()->setPrevNode(brother);
            }

            // 要删除的 key 为父节点中 node 和 brother 中间的键
            int nodeIndex, brotherIndex;
            for (int i = 0; i < father->getPointerCapacity(); i++) {
                if (father->getPointer(i) == node) {
                    nodeIndex = i;
                }
                if (father->getPointer(i) == brother) {
                    brotherIndex = i;
                }
            }
            delKey = father->getKey(min(nodeIndex, brotherIndex));

            delete node;
        } else {
            node->merge(brother);

            // 处理叶子节点关系，此时需要删除的是 node 的右兄弟
            node->setNextNode(brother->getNextNode());
            if (brother->getNextNode()) {
                brother->getNextNode()->setPrevNode(node);
            }

            // 要删除的 key 为父节点中 node 和 brother 中间的键
            int nodeIndex, brotherIndex;
            for (int i = 0; i < father->getPointerCapacity(); i++) {
                if (father->getPointer(i) == node) {
                    nodeIndex = i;
                }
                if (father->getPointer(i) == brother) {
                    brotherIndex = i;
                }
            }
            delKey = father->getKey(min(nodeIndex, brotherIndex));

            delete brother;
        }

        return this->removeInternalNode(father, delKey);
    }

    void print() {
        Node<K> *r = this->root;
        if (r == nullptr) {
            return;
        }

        queue<Node<K> *> queue;
        queue.push(r);

        int level = 0, count = 0;
        Node<K> *last, *nlast;
        last = r;
        nlast = r;
        cout << "Level " << ++level << endl;

        while (!queue.empty()) {
            Node<K> *node = queue.front();
            queue.pop();

            if (node->getType() == NodeType::LEAF) {
                cout << "| " << *((LeafNode<K, V>*) node) << " |" ;
                count++;
            } else {
                cout << "| " << *((InternalNode<K>*) node) << " |" ;
                count++;

                for (int i = 0; i < node->getLength() + 1; i++) {
                    Node<K> *next = ((InternalNode<K>*)node)->getPointer(i);
                    queue.push(next);
                    nlast = next;
                }
            }

            if (count % 2 == 0 || last == node) {
                cout << endl;
            }

            if (last == node) {
                last = nlast;
                count = 0;
                if (!queue.empty()) {
                    cout << "Level " << ++level << endl;
                }
            }
        }
    }
};

#endif //DBMS_BPLUSTREE_H
